#include"StringPool.h"
#include"variable.h"
#include<vector>
#include<ctime>
#include<iostream>
#include<map>
#include<iomanip>
#include"DependencyInstance.h"
#include"Util.h"
#include"KBestParseForest3O.h"
using namespace std;

void  KBestParseForest3O::getIndex(int g, int s, int m, int t, int *m_index, int *t_index){
	*m_index = m - s - 1;
	*t_index = t - m - 1;
}

int KBestParseForest3O::getKey(int dir, int k){
	return dir * K + k;
}

int KBestParseForest3O::getKey(int g, int s, int t, int dir, int comp, int k){
	int key = g;
	key = key * (end + 1) + s;
	key = key * (end + 1) + t;
	key = key * 2 + dir;
	key = key * 3 + comp;
	key = key * K + k;
	return key;
}

KBestParseForest3O::KBestParseForest3O(ParseForestItem *****d_chart, int start, int end, DependencyInstance &inst, int K) : KBestParseForest3OGSibling(start, end, inst, K){
	this->d_chart = d_chart;
	for(int i = 0; i < end + 1; i++){
		for(int j = 0; j< end + 1; j++){
			int nn1;
			if(j < i){
				nn1 = i - j - 1;
			}
			else if(i == j){
				nn1 = 0;
			}
			else{
				nn1 = end + 1 - j - 1;
			}
			for(int k = 0; k < nn1; k++){
				int k_real = k + j + 1;
				int nn2;
				if(j < i){
					nn2 = i - k_real - 1;
				}
				else{
					nn2 = end + 1 - k_real - 1;
				}
				for(int l = 0; l < nn2; l++){
					for(int dir = 0; dir < 2; dir++){
						for(int kk = 0; kk < K; kk++){
							int key = getKey(dir, kk);
							d_chart[i][j][k][l][key] = ParseForestItem(j, k, l, -1, dir, 3, Negative_Infinity, NULL, NULL);
						}
					}
				}
			}
		}
	}
}

KBestParseForest3O::~KBestParseForest3O(){}

bool KBestParseForest3O::add_d(int g, int s, int m, int t, int type, int dir, double score, ParseForestItem *p1, ParseForestItem *p2){
	bool added = false;
	int key = getKey(dir, K - 1);
	int m_index, t_index;
	getIndex(g, s, m, t, &m_index, &t_index);

	if(d_chart[g][s][m_index][t_index][key].prob > score){
		return false;
	}
	for(int i = 0;i< K; i++){
		key = getKey(dir, i);
		if(d_chart[g][s][m_index][t_index][key].prob < score){
			ParseForestItem tmp = d_chart[g][s][m_index][t_index][key];
			d_chart[g][s][m_index][t_index][key] = ParseForestItem(s, m, t, type, dir, 3, score, p1, p2);
			for(int j = i + 1; j < K && tmp.prob > Negative_Infinity; j++){
				int _key = getKey(dir, j);
				ParseForestItem tmp1 = d_chart[g][s][m_index][t_index][_key];
				d_chart[g][s][m_index][t_index][_key] = tmp;
				tmp = tmp1;
			}
			added = true;
			break;
		}
	}
	return added;
}

bool KBestParseForest3O::add_output_d(ParseForestItem *output_d, int m, int t, int type, double score, ParseForestItem *p1, ParseForestItem *p2){
	bool added = false;
	int key = getI_d(m, t, K - 1);

	if(output_d[key].prob > score){
		return false;
	}
	for(int i = 0; i < K; i++){
		key = getI_d(m, t, i);
		if(output_d[key].prob < score){
			ParseForestItem tmp = output_d[key];
			output_d[key] = ParseForestItem(0, m, t, type, 0, 3, score, p1, p2);
			for(int j = i + 1; j < K && tmp.prob > Negative_Infinity; j++){
				int _key = getI_d(m, t, j);
				ParseForestItem tmp1 = output_d[_key];
				output_d[_key] = tmp;
				tmp = tmp1;
			}
			added = true;
			break;
		}
	}
	return added;
}

ParseForestItem* KBestParseForest3O::getItems_d(int g, int s, int m, int t, int dir){
	int key = getKey(dir, 0);
	int m_index, t_index;
	getIndex(g, s, m, t, &m_index, &t_index);
	return (d_chart[g][s][m_index][t_index] + key);
}

double KBestParseForest3O::getProb_d(int g, int s, int m, int t, int dir, int i){
	int key = getKey(dir, i);
	int m_index, t_index;
	getIndex(g, s, m, t, &m_index, &t_index);
	return d_chart[g][s][m_index][t_index][key].prob;
}

double KBestParseForest3O::getProb_d(int g, int s, int m, int t, int dir){
	return getProb_d(g, s, m, t, dir, 0);
}

vector<pair<FeatureVector*, string*> >* KBestParseForest3O::getBestParses_strict(DependencyInstance *inst, DependencyPipe *pipe, std::vector<vector<vector<double> > > &probs, std::vector<vector<vector<double> > > &probs_trips, std::vector<vector<vector<double> > > &probs_sibs, std::vector<vector<vector<vector<double> > > > &probs_ttrips, std::vector<vector<vector<vector<double> > > > &probs_tsibs, std::vector<vector<vector<vector<double> > > > &nt_probs, bool labeled){
	vector<pair<FeatureVector*, string*> >* d = new vector<pair<FeatureVector*, string*> >(K);
	vector<ParseForestItem> result = vector<ParseForestItem>(K);
	for(int i = 0; i < K; i++){
		result[i] = ParseForestItem(0, -1, end, -1, 0, 1, Negative_Infinity, NULL, NULL);
	}
	for(int i = 1; i < end + 1; i++){
		ParseForestItem *left = getItems(0, 1, i, 1, 1);
		ParseForestItem *right = getItems(0, i, end, 0, 1);
		vector<BestPair>* pairs = getKBestPairs(left, right);
		int size = (int)(pairs->size());
		bool added = true;
		double prob_0i = probs[0][i][0];
		double nt_prob_s_01 = 0.0;
		double nt_prob_t_00 = 0.0;
		if(labeled){
			nt_prob_s_01 = nt_probs[0][rootType][0][1];
			nt_prob_t_00 = nt_probs[i][rootType][0][0];
			prob_0i += nt_prob_s_01 + nt_prob_t_00;
		}

		double prob_000i = probs_ttrips[0][0][0][i] + probs_tsibs[0][0][i][0] + probs_trips[0][0][i] + probs_sibs[0][i][0];

		for(int k = 0; k < size; k++){
			if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
				break;
			}

			int k1 = (*pairs)[k].first;
			int k2 = (*pairs)[k].second;
			double prob = left[k1].prob + right[k2].prob + prob_0i + prob_000i;
			added = false;
			for(int j = 0; j < K; j++){
				if(prob > result[j].prob){
					ParseForestItem tmp = result[j];
					result[j] = ParseForestItem(0, i, end, -1, 0, 1, prob, &(left[k1]), &(right[k2]));
					for(int m = j + 1; m < K; m++){
						ParseForestItem tmp1 = result[m];
						result[m] = tmp;
						tmp = tmp1;
					}
					added = true;
					break;
				}
			}
			if(!added){
				break;
			}
		}
		delete(pairs);
	}
	for(int i = 0; i < K; i++){
		if(result[i].prob > Negative_Infinity){

			string* str1 = getDepString(*(result[i].left));
			(*d)[i].second = Util::trim(str1);
			pool->push_back(str1);
			int ll = (int)((*d)[i].second->length());
			if(ll != 0){
				(*d)[i].second->append(" ");
			}
			(*d)[i].second->append("0|");
			str1 = Util::intToString(result[i].r);
			str1->append(":");
			(*d)[i].second->append(*str1);
			pool->push_back(str1);
			str1 = Util::intToString(KBestParseForest::rootType);
			str1->append(" ");
			(*d)[i].second->append(*str1);
			pool->push_back(str1);
			str1 = getDepString(*(result[i].right));
			string* str2 = Util::trim(str1);
			pool->push_back(str1);
			(*d)[i].second->append(*str2);
			pool->push_back(str2);
			str1 = (*d)[i].second;
			(*d)[i].second = Util::trim(str1);
			pool->push_back(str1);

			(*d)[i].first = getFeatureVector(inst, pipe, (*d)[i].second);
		}
		else{
			(*d)[i].first = NULL;
			(*d)[i].second = NULL;
		}
	}
	return d;
}

int KBestParseForest3O::getI_d(int i, int j, int k){
	int key = i;
	key = key * (end + 1) + j;
	key = key * K + k;
	return key;
}

vector<pair<FeatureVector*, string*> >* KBestParseForest3O::getBestParses_normal(DependencyInstance *inst, DependencyPipe *pipe, std::vector<vector<vector<double> > > &probs, std::vector<vector<vector<double> > > &probs_trips, std::vector<vector<vector<double> > > &probs_sibs, std::vector<vector<vector<vector<double> > > > &probs_ttrips, std::vector<vector<vector<vector<double> > > > &probs_tsibs, std::vector<vector<vector<vector<double> > > > &nt_probs, bool labeled){
	vector<pair<FeatureVector*, string*> >* d = new vector<pair<FeatureVector*, string*> >(K);
	ParseForestItem *output = new ParseForestItem[(end + 1) * 2 * K];
	for(int i = 0; i < end + 1; i++){
		for(int c = 0; c < 2; c++){
			for(int k = 0; k < K; k++){
				int key = getI(i, c, k);
				output[key] = ParseForestItem(0, -1, i, -1, 0, c, Negative_Infinity, NULL, NULL);
			}
		}
	}

	bool added = false;
	add_output(output, -1, 0, -1, 1, 0.0, NULL, NULL);

	ParseForestItem *output_d = new ParseForestItem[(end + 1) * ( end + 1) * K];
	for(int i = 0; i < end + 1; i++){
		for(int j = i + 1; j < end + 1; j++){
			for(int k = 0; k < K; k++){
				int key = getI_d(i, j, k);
				output_d[key] = ParseForestItem(0, i, j, -1, 0, 3, Negative_Infinity, NULL, NULL);
			}
		}
	}

	//int num = (end - 1) * end / 2;

	for(int t = 1; t < end + 1; t++){
		double prob_0t = probs[0][t][0];
		int type = KBestParseForest::rootType;
		double nt_prob_s_01 = 0.0;
		double nt_prob_t_00 = 0.0;
		if(labeled){
			nt_prob_s_01 = nt_probs[0][type][0][1];
			nt_prob_t_00 = nt_probs[t][type][0][0];
		}
		ParseForestItem* b1 = NULL;
		ParseForestItem* c1 = NULL;
		vector<BestPair>* pairs = NULL;
		int length1;
		int off;

		for(int m = 1; m < t; m++){
			b1 = getItems(0, 1, m, 1, 1);
			c1 = getItems(0, m, t, 1, 2);

			double prob_000m = probs_ttrips[0][0][0][m] + probs_tsibs[0][0][m][0] + probs_trips[0][0][m] + probs_sibs[0][m][0] + probs[0][m][0];
			
			double prob_00mt = probs_ttrips[0][0][m][t - m - 1] + probs_tsibs[0][m][t - m - 1][0] + probs_trips[0][m][t] + probs_sibs[m][t][1] + prob_0t;

			pairs = getKBestPairs(b1, c1);
			length1 = (int)(pairs->size());

			for(int k = 0; k < length1; k++){
				if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
					break;
				}

				int comp1 = (*pairs)[k].first;
				int comp2 = (*pairs)[k].second;

				double prob_fin = b1[comp1].prob + c1[comp2].prob + prob_000m + prob_00mt;
				if(labeled){
					prob_fin += nt_prob_s_01 + nt_prob_t_00;
				}
				added = add_output_d(output_d, m, t, type, prob_fin, &(b1[comp1]), &(c1[comp2]));
				if(!added){
					break;
				}
			}
			delete(pairs);

			for(int r = 1; r < m; r++){
				off = getI_d(r, m, 0);
				b1 = (output_d + off);
				c1 = getItems(0, m, t, 0, 2);
				
				pairs = getKBestPairs(b1, c1);
				length1 = (int)(pairs->size());

				double prob_0rmt = probs_ttrips[0][r][m - r - 1][t - m - 1] + probs_tsibs[r][m][t - m - 1][1] + probs_trips[0][m][t] + probs_sibs[m][t][1] + prob_0t;

				for(int k = 0; k < length1; k++){
					if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
						break;
					}

					int comp1 = (*pairs)[k].first;
					int comp2 = (*pairs)[k].second;

					double prob_fin = b1[comp1].prob + c1[comp2].prob + prob_0rmt;
					if(labeled){
						prob_fin += nt_prob_s_01 + nt_prob_t_00;
					}
					added = add_output_d(output_d, m, t, type, prob_fin, &(b1[comp1]), &(c1[comp2]));
					if(!added){
						break;
					}
				}
				delete(pairs);
			}
		}

		off = getI(0, 1, 0);
		b1 = (output + off);
		c1 = getItems(0, 1, t, 1, 1);
		double prob_000t = probs_ttrips[0][0][0][t] + probs_tsibs[0][0][t][0] + probs_trips[0][0][t] + probs_sibs[0][t][0] + prob_0t;
		pairs = getKBestPairs(b1, c1);
		length1 = (int)(pairs->size());
		for(int k = 0; k < length1; k++){
			if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
				break;
			}

			int comp1 = (*pairs)[k].first;
			int comp2 = (*pairs)[k].second;

			double prob_fin = b1[comp1].prob + c1[comp2].prob + prob_000t;
			if(labeled){
				prob_fin += nt_prob_s_01 + nt_prob_t_00;
			}
			added = add_output(output, 0, t, type, 0, prob_fin, &(b1[comp1]), &(c1[comp2]));
			if(!added){
				break;
			}
		}
		delete(pairs);

		for(int r = 1; r < t; r++){
			off = getI_d(r, t, 0);
			b1 = (output_d + off);
			c1 = getItems(0, t, t, 0, 1);
			pairs = getKBestPairs(b1, c1);
			length1 = (int)(pairs->size());

			for(int k = 0; k < length1; k++){
				if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
					break;
				}

				int comp1 = (*pairs)[k].first;
				int comp2 = (*pairs)[k].second;
				
				double bc = b1[comp1].prob + c1[comp2].prob;
				
				added = add_output(output, r, t, type, 0, bc, &(b1[comp1]), &(c1[comp2]));
				if(!added){
					break;
				}
			}
			delete(pairs);
		}

		for(int r = 1; r <= t; r++){
			off = getI(r, 0, 0);
			b1 = (output + off);
			c1 = getItems(0, r, t, 0, 1);
			pairs = getKBestPairs(b1, c1);
			length1 = (int)(pairs->size());
			for(int k = 0; k < length1; k++){
				if(((*pairs)[k]).first == -1 || ((*pairs)[k]).second == -1){
					break;
				}

				int comp1 = (*pairs)[k].first;
				int comp2 = (*pairs)[k].second;

				double prob_fin = b1[comp1].prob + c1[comp2].prob;
				added = add_output(output, r, t, -1, 1, prob_fin, &(b1[comp1]), &(c1[comp2]));
				if(!added){
					break;
				}
			}
			delete(pairs);
		}
	}

	for(int i = 0; i < K; i++){
		int key = getI(end, 1, i);
		if(output[key].prob > Negative_Infinity){
			(*d)[i].second = getDepString(output[key]);
			(*d)[i].first = getFeatureVector(inst, pipe, (*d)[i].second);
		}
		else{
			(*d)[i].first = NULL;
			(*d)[i].second = NULL;
		}
	}
	delete[](output);
	delete[](output_d);
	return d;
}

vector<pair<FeatureVector*, string*> >* KBestParseForest3O::getBestParses(DependencyInstance *inst, DependencyPipe *pipe, std::vector<vector<vector<double> > > &probs, std::vector<vector<vector<double> > > &probs_trips, std::vector<vector<vector<double> > > &probs_sibs, std::vector<vector<vector<vector<double> > > > &probs_ttrips, std::vector<vector<vector<vector<double> > > > &probs_tsibs, std::vector<vector<vector<vector<double> > > > &nt_probs, bool labeled, bool normal){
	if(normal){
		return getBestParses_normal(inst, pipe, probs, probs_trips, probs_sibs, probs_ttrips, probs_tsibs, nt_probs, labeled);
	}
	else{
		return getBestParses_strict(inst, pipe, probs, probs_trips, probs_sibs, probs_ttrips, probs_tsibs, nt_probs, labeled);
	}
}

string* KBestParseForest3O::getDepString(ParseForestItem &pfi){
	if(pfi.left == NULL){
		return pool->get();
	}

	if(pfi.dir == 0 && pfi.comp == 0){
		string* str1 = getDepString(*(pfi.left));
		str1->append(" ");
		string* str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		string* result = Util::trim(str1);
		pool->push_back(str1);
		pool->push_back(str2);
		result->append(" ");
		str1 = Util::intToString(pfi.s);
		result->append(*str1 + "|");
		pool->push_back(str1);
		str1 = Util::intToString(pfi.t);
		result->append(*str1 + ":");
		pool->push_back(str1);
		str1 = Util::intToString(pfi.type);
		result->append(*str1);
		pool->push_back(str1);
		str1 = result;
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else if(pfi.dir == 1 && pfi.comp == 0){
		string* str1 = Util::intToString(pfi.t);
		str1->append("|");
		string* str2 = Util::intToString(pfi.s);
		str2->append(":");
		str1->append(*str2);
		pool->push_back(str2);
		str2 = Util::intToString(pfi.type);
		str2->append(" ");
		str1->append(*str2);
		pool->push_back(str2);
		string* result = str1;
		str1 = getDepString(*(pfi.left));
		str1->append(" ");
		str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		pool->push_back(str2);
		str2 = Util::trim(str1);
		pool->push_back(str1);
		result->append(*str2);
		pool->push_back(str2);
		str1 = result;
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else if(pfi.dir == 0 && pfi.comp == 3){
		string* str1 = getDepString(*(pfi.left));
		string *result = Util::trim(str1);
		pool->push_back(str1);
		int ll = (int)(result->length());
		if(ll != 0){
			result->append(" ");
		}
		str1 = Util::intToString(pfi.s);
		result->append(*str1);
		pool->push_back(str1);
		result->append("|");
		str1 = Util::intToString(pfi.r);
		result->append(*str1);
		pool->push_back(str1);
		result->append(":");
		str1 = Util::intToString(pfi.type);
		result->append(*str1);
		pool->push_back(str1);
		result->append(" ");
		str1 = getDepString(*(pfi.right));
		string *str2 = Util::trim(str1);
		pool->push_back(str1);
		str1 = result;
		str1->append(*str2);
		pool->push_back(str2);
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else if(pfi.dir == 1 && pfi.comp == 3){
		string* str1 = getDepString(*(pfi.left));
		string *result = Util::trim(str1);
		pool->push_back(str1);
		int ll = (int)(result->length());
		if(ll != 0){
			result->append(" ");
		}
		str1 = Util::intToString(pfi.t);
		result->append(*str1);
		pool->push_back(str1);
		result->append("|");
		str1 = Util::intToString(pfi.r);
		result->append(*str1);
		pool->push_back(str1);
		result->append(":");
		str1 = Util::intToString(pfi.type);
		result->append(*str1);
		pool->push_back(str1);
		result->append(" ");
		str1 = getDepString(*(pfi.right));
		string *str2 = Util::trim(str1);
		pool->push_back(str1);
		str1 = result;
		str1->append(*str2);
		pool->push_back(str2);
		result = Util::trim(str1);
		pool->push_back(str1);
		return result;
	}
	else{
		string* str1 = getDepString(*(pfi.left));
		str1->append(" ");
		string* str2 = getDepString(*(pfi.right));
		str1->append(*str2);
		string* result = Util::trim(str1);
		pool->push_back(str1);
		pool->push_back(str2);
		return result;
	}
}
